from pathlib import Path

# from rapidocr import RapidOCR, VisRes
from rapid_table import RapidTable, RapidTableInput, VisTable

# 开启onnx-gpu推理
# input_args = RapidTableInput(use_cuda=True)
# table_engine = RapidTable(input_args)

# 使用torch推理版本的unitable模型
# input_args = RapidTableInput(model_type="unitable", use_cuda=True, device="cuda:0")
# table_engine = RapidTable(input_args)

# ocr_engine = RapidOCR()
# vis_ocr = VisRes()

# 默认是slanet_plus模型
input_args = RapidTableInput(model_type="unitable")
table_engine = RapidTable(input_args)
viser = VisTable()

# img_paths = ["/home/fengjie/doc-parser/MinerU/src/RapidTable/40fb854e57b5258d90a73c4ad3ad52cf.png","/home/fengjie/doc-parser/MinerU/src/RapidTable/3a417bf2cc6a33e21bf2cd995da6898bd1de20d6a35b24157b27a77244616b1d.jpg"]
img_paths = ["/home/fengjie/doc-parser/MinerU/src/TableStructureRec/18b104d898cfb5ea82b04f8bb39e62b3.png"]
pipeline_yaml="../config_yaml/OCR.yaml"

# OCR
# rapid_ocr_output = ocr_engine(img_path, return_word_box=True)
# ocr_result = list(
# zip(rapid_ocr_output.boxes, rapid_ocr_output.txts, rapid_ocr_output.scores)
# )
import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from paddlex_table.ocr import paddlex_ocr

device = "gpu"  # or "cpu"
ocr_results = paddlex_ocr(img_paths, pipeline_yaml=pipeline_yaml, device=device)

for i, img_path in enumerate(img_paths):
    ocr_result = ocr_results[i]
    ocr_result = list(zip(ocr_result['rec_polys'],ocr_result['rec_texts'],ocr_result['rec_scores']))
        

    table_results = table_engine(img_path, ocr_result)
    table_html_str, table_cell_bboxes = table_results.pred_html, table_results.cell_bboxes
    # Save
    save_dir = Path("outputs")
    save_dir.mkdir(parents=True, exist_ok=True)

    save_html_path = save_dir / f"{Path(img_path).stem}.html"
    save_drawed_path = save_dir / f"{Path(img_path).stem}_table_vis{Path(img_path).suffix}"
    save_logic_points_path = save_dir / f"{Path(img_path).stem}_table_col_row_vis{Path(img_path).suffix}"

    # 可视化表格recc结果
    vis_imged = viser(img_path, table_results, save_html_path, save_drawed_path, save_logic_points_path)

    print(f"The results has been saved {save_dir}")